package com.tsx.test.utils;

import com.tsx.test.entity.dto.DeptNodeDTO;
import com.tsx.test.entity.po.Department;
import com.tsx.test.entity.vo.DeptNode;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

public class TreeUtil {
    public static List<DeptNode> buildTree(List<DeptNode> nodes, Long parentId) {
        List<DeptNode> tree = new ArrayList<>();
        for (DeptNode node : nodes) {
            if (Objects.equals(node.getParentId(), parentId)) {
                tree.add(node);
            }
        }
        for (DeptNode node : tree) {
            buildChildren(node, nodes);
        }
        return tree;
    }

    private static void buildChildren(DeptNode node, List<DeptNode> nodes) {
        for (DeptNode child : nodes) {
            if (child.getParentId().equals(node.getId())) {
                node.getChildren().add(child);
                buildChildren(child, nodes);
            }
        }
    }

    public static List<DeptNode> buildTree2(List<DeptNode> nodes, Long parentId) {
        Map<Long, DeptNode> idMap = nodes.stream()
                .collect(Collectors.toMap(DeptNode::getId, node -> node));
        List<DeptNode> tree = new ArrayList<>();
        for (DeptNode node : nodes) {
            if (Objects.equals(node.getParentId(), parentId)) {
                tree.add(node);
            } else {
                DeptNode parentNode = idMap.get(node.getParentId());
                if (parentNode != null) {
                    parentNode.getChildren().add(node);
                }
            }
        }
        return tree;
    }

    public static List<Department> flattenTree(List<DeptNodeDTO> nodes, Long parentId) {
        List<Department> result = new ArrayList<>();
        for (DeptNodeDTO node : nodes) {
            Department dept = new Department();
            dept.setId(node.getId());
            dept.setName(node.getName());
            dept.setLeader(node.getLeader());
            dept.setParentId(parentId);
            result.add(dept);

            if (node.getChildren() != null && !node.getChildren().isEmpty()) {
                List<Department> children = flattenTree(node.getChildren(), node.getId());
                result.addAll(children);
            }
        }
        return result;
    }

}
